// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/common/base/exceptions.h>
#include <pollux/common/base/simd_util.h>
#include <pollux/vector/flat_vector.h>
#include <pollux/vector/simple_vector.h>
#include <pollux/vector/type_aliases.h>

namespace kumo::pollux {
    /**
     * Provides a vector type that enacts run-length encoding for vectors that have
     * repeated runs of the same value.  For example, this would be useful for
     * {foo, foo, foo, foo, bar, bar, bar, foo, foo, foo, foo, foo} as it can be
     * reduced to {foo, bar, foo} along with run length information of {4, 3, 5}.
     *
     * NOTE - current implementation is not thread safe when accessing values.
     */
    template<typename T>
    class SequenceVector : public SimpleVector<T> {
    public:
        static constexpr bool can_simd =
        (std::is_same_v<T, int64_t> || std::is_same_v<T, int32_t> ||
         std::is_same_v<T, int16_t> || std::is_same_v<T, int8_t> ||
         std::is_same_v<T, size_t>);

        SequenceVector(
            pollux::memory::MemoryPool *pool,
            size_t length,
            VectorPtr sequenceValues,
            BufferPtr sequenceLengths,
            const SimpleVectorStats<T> &stats = {},
            std::optional<int32_t> distinctCount = std::nullopt,
            std::optional<vector_size_t> nullCount = std::nullopt,
            std::optional<bool> sorted = std::nullopt,
            std::optional<ByteCount> represented_bytes = std::nullopt,
            std::optional<ByteCount> storageByteCount = std::nullopt);

    public:
        ~SequenceVector() override = default;

        bool may_have_nulls() const override {
            return sequenceValues_->may_have_nulls();
        }

        bool may_have_nulls_recursive() const override {
            return sequenceValues_->may_have_nulls_recursive();
        }

        bool isNullAtFast(vector_size_t idx) const;

        bool is_null_at(vector_size_t idx) const override {
            return isNullAtFast(idx);
        }

        bool contains_null_at(vector_size_t idx) const override {
            if constexpr (std::is_same_v<T, ComplexType>) {
                if (is_null_at(idx)) {
                    return true;
                }

                size_t offset = offsetOfIndex(idx);
                return sequenceValues_->contains_null_at(offset);
            } else {
                return is_null_at(idx);
            }
        }

        const T valueAtFast(vector_size_t idx) const;

        const T value_at(vector_size_t idx) const override {
            return valueAtFast(idx);
        }

        std::unique_ptr<SimpleVector<uint64_t> > hash_all() const override;

        /**
         * Loads a 256bit vector of data at the virtual byteOffset given
         * Note this method is implemented on each vector type, but is intentionally
         * not virtual for performance reasons
         *
         * @param byteOffset - the byte offset to laod from
         */
        ksimd::batch<T> loadSIMDValueBufferAt(size_t index) const;

        /**
         * Returns a shared_ptr to the underlying byte buffer holding the values for
         * this vector as SequenceLength types. This is used during execution to
         * process over the subset of values when possible.
         */
        inline VectorPtr getSequenceValues() const {
            return sequenceValues_;
        }

        BaseVector *loaded_vector() override {
            auto loaded = BaseVector::loaded_vector_shared(sequenceValues_);
            if (loaded == sequenceValues_) {
                return this;
            }
            sequenceValues_ = loaded;
            setInternalState();
            return this;
        }

        const BaseVector *loaded_vector() const override {
            return const_cast<SequenceVector<T> *>(this)->loaded_vector();
        }

        const VectorPtr &value_vector() const override {
            return sequenceValues_;
        }

        BufferPtr getSequenceLengths() const {
            return sequenceLengths_;
        }

        /**
         * Returns a shared_ptr to the underlying values container.
         */
        inline VectorPtr getValuesVectorShared() const {
            return sequenceValues_;
        }

        vector_size_t numSequences() const {
            return sequenceValues_->size();
        }

        const BufferPtr &wrap_info() const override {
            return sequenceLengths_;
        }

        uint64_t retained_size() const override {
            return sequenceValues_->retained_size() + sequenceLengths_->capacity();
        }

        bool is_scalar() const override {
            return sequenceValues_->is_scalar();
        }

        /**
         * @returns the physical index for the location of the sequence value that
         * is logically resident at the given index, or, -1 for an index that is < 0
         * or an index that is >= the logical length of this vector.
         */
        vector_size_t offsetOfIndex(vector_size_t idx) const;

        const BaseVector *wrapped_vector() const override {
            return sequenceValues_->wrapped_vector();
        }

        vector_size_t wrapped_index(vector_size_t index) const override {
            return sequenceValues_->wrapped_index(offsetOfIndex(index));
        }

        void add_nulls(const uint64_t *bits, const SelectivityVector &rows) override {
            throw std::runtime_error("add_nulls not supported");
        }

        void add_nulls(const SelectivityVector &rows) override {
            throw std::runtime_error("add_nulls not supported");
        }

        std::string toString(vector_size_t index) const override {
            if (BaseVector::is_null_at(index)) {
                return "null";
            }
            auto inner = offsetOfIndex(index);
            std::stringstream out;
            out << "[" << index << "->" << inner << "] "
                    << sequenceValues_->toString(inner);
            return out.str();
        }

        VectorPtr slice(vector_size_t, vector_size_t) const override {
            POLLUX_NYI();
        }

        bool is_nulls_writable() const override {
            return false;
        }

        VectorPtr copy_preserve_encodings(
            pollux::memory::MemoryPool *pool = nullptr) const override {
            auto selfPool = pool ? pool : BaseVector::pool_;
            return std::make_shared<SequenceVector<T> >(
                selfPool,
                BaseVector::length_,
                sequenceValues_->copy_preserve_encodings(),
                AlignedBuffer::copy(selfPool, sequenceLengths_),
                SimpleVector<T>::stats_,
                BaseVector::distinctValueCount_,
                BaseVector::nullCount_,
                SimpleVector<T>::isSorted_,
                BaseVector::representedByteCount_,
                BaseVector::storageByteCount_);
        }

    private:
        // Prepares for use after construction.
        void setInternalState();

        bool checkLoadRange(size_t idx, size_t count) const;

        VectorPtr sequenceValues_;
        SimpleVector<T> *scalarSequenceValues_ = nullptr;
        BufferPtr sequenceLengths_;
        // Caches 'sequenceLengths_->as<vector_size_t>()'
        const vector_size_t *lengths_ = nullptr;

        // manage the state of the offset index operations to allow for the common
        // path of iterating in order to not require repeated trips through the
        // sequenceLengths_ array
        mutable vector_size_t lastIndexRangeStart_ = 0;
        mutable vector_size_t lastIndexRangeEnd_ = 0;
        mutable vector_size_t lastRangeIndex_ = 0;
    };

    template<typename T>
    using SequenceVectorPtr = std::shared_ptr<SequenceVector<T> >;
} // namespace kumo::pollux

#include <pollux/vector/sequence_vector-inl.h>
